"""
添加文档
"""
import os
import sys

path = os.path.dirname(__file__)
sys.path.append(os.path.join(path, '..'))

import math
import random
from datetime import datetime
from faker import Faker
from conf.mongo_conf import BASIC_DATA_DB
from conf.config import COMPANY_STAFFS_LIST, COMPANY_STAFFS_MIN_LIST


class InsertDocument(object):
    """
    添加文档
    """
    fake = Faker(locale='zh_CN')
    company_staffs = BASIC_DATA_DB[COMPANY_STAFFS_LIST]
    company_staffs_min = BASIC_DATA_DB[COMPANY_STAFFS_MIN_LIST]

    def data(self):
        sex_num = random.randint(0, 1)
        sex = bool(sex_num)
        document = {
            'name': self.fake.name_female() if sex else self.fake.name_male(),  # 员工姓名
            'sex': sex_num,  # 员工性别
            'age': random.randint(20, 55),  # 员工年龄
            'phone_number': self.fake.phone_number(),  # 员工手机号
            'company': self.fake.company(),  # 员工公司名
            'company_email': self.fake.company_email(),  # 员工公司邮箱
            'job': self.fake.job(),  # 员工职位
            'card_number': self.fake.credit_card_number(),  # 员工信用卡号
            'id_card': self.fake.ssn(),  # 员工身份证号
            'native_place_province': self.fake.province(),  # 员工籍贯省份
            'native_place_address': self.fake.address(),  # 员工籍贯详细地址
            'create_time': datetime.utcnow()  # 创建时间
        }
        return document

    def insert_one(self):
        """单个添加"""
        self.company_staffs.insert_one(self.data())

    def insert_many(self):
        """批量添加40万数据"""
        for_times = math.ceil(400000/1000)
        for i in range(for_times):
            documents = []
            for y in range(1000):
                documents.append(self.data())
            self.company_staffs.insert_many(documents)
            print(f"""第{i}次，添加1000条数据成功""")

    def insert_many_min(self):
        documents = []
        for _ in range(1000):
            document = self.data()
            documents.append({
                'name': document.get('name'),  # 员工姓名
                'sex': document.get('sex'),  # 员工性别
                'age': document.get('age'),  # 员工年龄
                'phone_number': document.get('phone_number'),  # 员工手机号
            })
        self.company_staffs_min.insert_many(documents)

    def run(self):
        # self.insert_one()   # 添加一条数据
        #         # self.insert_many()   # 批量添加多条数据
        self.insert_many_min()   # 批量添加多条数据


if __name__ == '__main__':
    insert_document = InsertDocument()
    insert_document.run()
